from fastapi import FastAPI
from database.milvus_client import MilvusClient
from routers.chat_router import router as chat_router
import uvicorn
import os
from dotenv import load_dotenv

from server.chat_record_server import ChatRecordServer

from routers.memory_router import router as memory_router
from routers.image_text_router import router as image_router
from routers.encryption_router import router as encryption_router
from server.image_text_vac_server import ImageTextVacServer
from server.memory_record_server import MemoryRecordServer
from config.common_config import DATABASE_CONFIG
# 加载环境变量
# load_dotenv()

app = FastAPI(
    title="Enhanced Chat Service API",
    description="支持多内容类型和角色的聊天记录服务",
    version="2.0.0"
)

# 包含路由
app.include_router(chat_router)
app.include_router(memory_router)
app.include_router(image_router)
app.include_router(encryption_router)

# 初始化 Milvus 数据库

chat_record_server = ChatRecordServer(
    host=DATABASE_CONFIG["HOST"],
    port=DATABASE_CONFIG["PORT"],
    user=DATABASE_CONFIG["USER"],
    pwd=DATABASE_CONFIG["PASSWORD"]
)

memory_record_server = MemoryRecordServer(
    host=DATABASE_CONFIG["HOST"],
    port=DATABASE_CONFIG["PORT"],
    user=DATABASE_CONFIG["USER"],
    pwd=DATABASE_CONFIG["PASSWORD"]
                                              )


image_text_server= ImageTextVacServer(
    host=DATABASE_CONFIG["HOST"],
    port=DATABASE_CONFIG["PORT"],
    user=DATABASE_CONFIG["USER"],
    pwd=DATABASE_CONFIG["PASSWORD"]
)


@app.on_event("startup")
async def startup_event():
    chat_record_server.connect()
    chat_record_server.init_collection()
    chat_record_server.load_collection()
    memory_record_server.connect()
    memory_record_server.init_collection()
    memory_record_server.load_collection()



    image_text_server.connect()
    image_text_server.init_collection()
    image_text_server.load_collection()

    print("Service started and Milvus collectio loaded")

@app.on_event("shutdown")
async def shutdown_event():
    chat_record_server.release_collection()
    chat_record_server.disconnect()
    memory_record_server.release_collection()
    memory_record_server.disconnect()

    image_text_server.release_collection()
    image_text_server.disconnect()
    print("Service stopped and Milvus connection released")
    print(f"")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=10089)



